import os
import math
import torch
import numpy as np
from torch import nn
import numpy.typing as npt
from einops import rearrange
from typing import Optional, Iterable, BinaryIO, IO


def cross_entropy(logits: torch.Tensor, target: torch.Tensor):
    """
    计算一个 batch 中所有 token 的平均交叉熵

    Args:
        logits (torch.Tensor):  未归一化的 logits 向量，形状为 (batch_size, ..., vocab_size)
        target (torch.Tensor):  表示 logits 中第几个是正确答案，形状为 (batch_size, ...)

    Returns:
        torch.Tensor: 平均交叉熵
    """

    # 原式 = -log{softmax(logits)[target]}
    #   得 = -log{exp(logits[target])} - (-log{sum(exp(logits))})
    # 化简 = -logits[target] + log(sum(exp(logits)))

    # 将 logits 的 batch、seq_len 等维度合并
    # (batch_size*seq_len, vocab_size)
    logits = rearrange(logits, "b ... v -> (b ...) v")
    # target 也合并
    # (batch_size*seq_len)
    target = rearrange(target, "b ... -> (b ...)")

    # 将每个 logits 向量减去当前行的最大值，防止向上溢出
    logits = logits - logits.max(dim=-1, keepdim=True).values

    # 计算 logits[target]
    logits_target = logits.gather(1, target.unsqueeze(1)).squeeze(1)
    # 计算 log(sum(exp(logits)))
    log_sum_exp = torch.log(torch.sum(torch.exp(logits), dim=-1))

    # 交叉熵即为 -logits[target] + log(sum(exp(logits)))
    cross_entropy_loss = -logits_target + log_sum_exp

    # 求平均值
    return cross_entropy_loss.mean()


def lr_cosine_schedule(t: int, a_max: float, a_min: float, t_w: int, t_c: int):
    """
    余弦退火学习率

    Args:
        t (int):        当前的迭代次数
        a_max (float):  最大学习率
        a_min (float):  最小学习率
        t_w (int):      线性预热的迭代数
        t_c (int):      余弦退火的迭代数
    """
    if t < t_w:
        return t / t_w * a_max
    elif t <= t_c:
        theta = (t - t_w) / (t_c - t_w) * math.pi
        return a_min + 0.5 * (1 + math.cos(theta)) * (a_max - a_min)
    else:
        return a_min


def gradient_clipping(
    parameters: Iterable[nn.Parameter], max_l2_norm: float, eps: float = 1e-6
):
    """
    梯度裁剪

    Args:
        parameters (Iterable[nn.Parameter]):    所有参数
        max_l2_norm (float):                    L2范数的裁剪阈值
        eps (float):                            误差偏移量
    """

    # 求所有参数梯度的 L2 范数
    total_norm = 0.0
    for p in parameters:
        if p.grad is None:
            continue
        total_norm += torch.sum(p.grad.data**2)
    total_norm = total_norm**0.5

    # 若超过裁剪阈值，则进行梯度裁剪
    if total_norm > max_l2_norm:
        factor = max_l2_norm / (total_norm + eps)

        for p in parameters:
            if p.grad is None:
                continue
            p.grad.data *= factor


def get_batch(
    dataset: npt.NDArray,
    batch_size: int,
    context_length: int,
    device: str | torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    传入一个 token 序列数据集，采样 batch 条长度为 context_length 的 token 序列数据

    Args:
        dataset (npt.NDArray):          token 序列数据集，形状为一个一维数组
        batch_size (int):               每个 batch 的大小
        context_length (int):           上下文长度，即 batch 中每个样本的序列长度
        device (str | torch.device):    保存采样数据的设备

    Returns:
        tuple[torch.Tensor, torch.Tensor]: 两个 Tensor.
            - 输入序列，形状为 (batch_size, context_length)
            - 目标序列，形状为 (batch_size, context_length)，每个输入序列向右平移一个 token 所得的序列。
    """
    dataset_length = dataset.shape[0]
    if dataset_length < context_length:
        raise ValueError(
            f"Dataset length {dataset_length} is less than context_length {context_length}."
        )

    starts = np.random.randint(0, dataset_length - context_length, batch_size)
    inputs = np.stack([dataset[s : s + context_length] for s in starts])
    targets = np.stack([dataset[s + 1 : s + context_length + 1] for s in starts])

    return torch.from_numpy(inputs).to(device), torch.from_numpy(targets).to(device)


def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    iteration: int,
    out: str | os.PathLike | BinaryIO | IO[bytes],
):
    """
    保存检查点。将模型参数、优化器状态、迭代次数存入 out 路径的文件中。

    Args:
        model (nn.Module):                              模型对象
        optimizer (torch.optim.Optimizer):              优化器对象
        iteration (int):                                迭代次数
        out (str | os.PathLike | BinaryIO | IO[bytes]): checkpoint 保存路径
    """

    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "iteration": iteration,
        },
        out,
    )


def load_checkpoint(
    src: str | os.PathLike | BinaryIO | IO[bytes],
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer],
) -> int:
    """
    加载检查点。从 src 路径的文件中加载模型参数、优化器状态，并返回迭代次数。

    Args:
        src (str | os.PathLike | BinaryIO | IO[bytes]): checkpoint 保存路径
        model (nn.Module):                              模型对象
        optimizer (torch.optim.Optimizer):              优化器对象
    """

    obj = torch.load(src)
    model.load_state_dict(obj["model_state_dict"])
    if optimizer is not None:
        optimizer.load_state_dict(obj["optimizer_state_dict"])
    return obj["iteration"]


def evaluate_model(
    model: nn.Module,
    validate_dataset: npt.NDArray,
    batch_size: int,
    context_length: int,
    device: str | torch.device,
    eval_times: int = 10,
):
    """
    在验证集上对模型性能进行评估。

    Args:
        model (nn.Module):              模型对象
        validate_dataset (npt.NDArray): 验证集，形状为一个一维的数组
        batch_size (int):               每个 batch 的大小
        context_length (int):           上下文长度，即 batch 中每个样本的序列长度
        device (str | torch.device):    保存采样数据的设备
        eval_times (int, optional):     评估的次数
    Returns:
        float:  表示 eval_times 次评估结果的平均值
    """

    model.eval()  # 切换为 evaluate 模式

    total_loss = 0.0
    with torch.no_grad():  # 禁用梯度相关的计算与存储
        for _ in range(eval_times):
            inputs, targets = get_batch(
                dataset=validate_dataset,
                batch_size=batch_size,
                context_length=context_length,
                device=device,
            )

            logits = model(inputs)  # 前向传播
            loss = cross_entropy(logits, targets)  # 计算损失

            total_loss += loss

    model.train()  # 切换回 train 模式

    return total_loss / eval_times


def compute_grad_norm(parameters: Iterable[nn.Parameter]):
    """
    计算参数梯度的 L2 范数

    参数:
        parameters: 模型参数

    返回:
        梯度的 L2 范数
    """
    total_norm = 0.0
    for p in parameters:
        if p.grad is not None:
            param_norm = p.grad.detach().data.norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm**0.5
    return total_norm
